/***********************************************************************************
 * ELEC5616 Computer and Network Security, The University of Sydney
 * 
 * PROJECT: StealthNet FILENAME: StealthNetComms.java AUTHORS: Stephen Gould, Matt Barrie, Ryan Junee DESCRIPTION:
 * Implementation of StealthNet Communications for ELEC5616 programming assignment. This code has been written for the
 * purposes of teaching cryptography and computer security. It is to be used as a demonstration only. No attempt has
 * been made to optimise the source code. VERSION: 1.0 IMPLEMENTS: initiateSession(); acceptSession();
 * terminateSession(); sendPacket(); recvPacket(); recvReady();
 * 
 * REVISION HISTORY:
 * 
 **********************************************************************************/

/* Import Libraries ********************************************************* */

import java.io.BufferedReader;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.FileWriter;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.PrintWriter;
import java.net.Socket;
import java.security.AlgorithmParameterGenerator;
import java.security.AlgorithmParameters;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.Signature;
import java.security.spec.X509EncodedKeySpec;
import java.util.Date;
import java.util.Scanner;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.CipherInputStream;
import javax.crypto.CipherOutputStream;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.KeyAgreement;
import javax.crypto.Mac;
import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.interfaces.DHPublicKey;
import javax.crypto.spec.DHParameterSpec;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.PBEParameterSpec;
import javax.crypto.spec.SecretKeySpec;

import sun.misc.IOUtils;

import com.sun.org.apache.xml.internal.security.exceptions.Base64DecodingException;
import com.sun.org.apache.xml.internal.security.utils.Base64;

/* StealthNetComms class **************************************************** */

public class StealthNetComms {
    public static final String SERVERNAME = "localhost";
    public static final int SERVERPORT = 5618;

    private Socket commsSocket; // communications socket
    private PrintWriter dataOut; // output data stream
    private BufferedReader dataIn; // input data stream

    private PublicKey dhPublicKey; // received public key
    private KeyPair dhKeyPair; // local generated key pair for Diffie-hellman
    private KeyAgreement keyAgree; // key agreement for the session
    private Cipher encrypt; // cipher used for the session
    private Cipher decrypt; // cipher used for the session
    private Mac hmac; // hmac algorithm used for the session
    private PublicKey rsaPublicKey;
    private PrivateKey rsaPrivateKey;
    private Signature signSignature;
    private Signature verifySignature;
    private String self;
    private String userId;

    private static final char[] HEXTABLE = { '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D',
        'E', 'F' };

    private static final long TOLERANCE = 5 * 60 * 1000; // 5 mins expiry time
                                                         // for a packet
    private static final String PRIVATE_KEY_EXTENTION = ".pri";
    private static final String PUBLIC_KEY_EXTENTION = ".pub";
    private static final String SHA1_MAC_EXTENTION = ".sha1";
    private static final byte[] SALT = "passw0rd".getBytes();

    private static final String FILE_AES_KEY = "shf9842yroh4fhyh43wr09843rt03";
    private static final String FILE_MAC_KEY = "fwei9ijf982h82j283j239h82j23f";
    
    public StealthNetComms() {
        commsSocket = null;
        dataIn = null;
        dataOut = null;
        dhPublicKey = null;
        dhKeyPair = null;
        keyAgree = null;
        encrypt = null;
        decrypt = null;
        hmac = null;
        rsaPublicKey = null;
        rsaPrivateKey = null;
        signSignature = null;
        verifySignature = null;
        self = null;
        userId = null;
    }
    
//    public void setUserId(String userId) throws Exception {
//        this.userId = userId;
//        this.rsaPublicKey = getRSAPublicKey(new File(userId + PUBLIC_KEY_EXTENTION));
//    }
//    
    protected void finalize() throws IOException {
        if (dataOut != null) dataOut.close();
        if (dataIn != null) dataIn.close();
        if (commsSocket != null) commsSocket.close();

        dhPublicKey = null;
        dhKeyPair = null;
        keyAgree = null;
        encrypt = null;
        decrypt = null;
        hmac = null;
        rsaPublicKey = null;
        rsaPrivateKey = null;
        signSignature = null;
        verifySignature = null;
        self = null;
        userId = null;
    }

    public void initiateSession(Socket socket, String self, String userId) {
        try {
            commsSocket = socket;
            dataOut = new PrintWriter(commsSocket.getOutputStream(), true);
            dataIn = new BufferedReader(new InputStreamReader(commsSocket.getInputStream()));
            this.self = self;
            this.userId = userId;
            initRSAKeyPair();
//            initSignature();
            initDHKey();
            recvDHPublicKey();
            initCipher();
            initMac();
        } catch (Exception e) {
            e.printStackTrace();
            System.err.println("Connection terminated.");
            System.exit(1);
        }
    }

    public void acceptSession(Socket socket, String self, String userId) {
        try {
            commsSocket = socket;
            dataOut = new PrintWriter(commsSocket.getOutputStream(), true);
            dataIn = new BufferedReader(new InputStreamReader(commsSocket.getInputStream()));
//            initRSAKeyPair();
//            initSignature();
            recvDHPublicKey();
            initDHKey();
            initCipher();
            initMac();
        } catch (Exception e) {
            e.printStackTrace();
            System.err.println("Connection terminated.");
            System.exit(1);
        }
    }

    public boolean terminateSession() {
        try {
            if (commsSocket == null) return false;
            dataIn.close();
            dataOut.close();
            commsSocket.close();
            commsSocket = null;

            dhPublicKey = null;
            dhKeyPair = null;
            keyAgree = null;
            encrypt = null;
            decrypt = null;
            hmac = null;
            rsaPublicKey = null;
            rsaPrivateKey = null;
            signSignature = null;
            verifySignature = null;
            self = null;
            userId = null;

        } catch (Exception e) {
            return false;
        }

        return true;
    }

    public boolean sendPacket(byte command) {
        return sendPacket(command, new byte[0]);
    }

    public boolean sendPacket(byte command, String data) {
        System.out.println("String data: " + data);
        return sendPacket(command, data.getBytes());
    }

    public boolean sendPacket(byte command, byte[] data) {
        return sendPacket(command, data, data.length);
    }

    public boolean sendPacket(byte command, byte[] data, int size) {
        StealthNetPacket pckt = new StealthNetPacket();
        pckt.command = command;
        pckt.data = new byte[size];
        System.arraycopy(data, 0, pckt.data, 0, size);
        return sendPacket(pckt);
    }

    public boolean sendPacket(StealthNetPacket pckt) {
        if (dataOut == null) return false;

        // if channel has been secured and a MAC mechanism set up
        // combine the message with a MAC code for message integrity
        if (hmac != null)
        // calculate mac code of the message (command + data)
            pckt.mac = hmac.doFinal(concatenate(pckt.command, pckt.data));

        System.out.println("[sent] " + pckt.toString());

        // signSignature.update(pckt.toString().getBytes());
        // byte[] sign = signSignature.sign();

        if (encrypt != null)
            try {
                dataOut.println(Base64.encode(encrypt.doFinal(pckt.toString().getBytes()), Integer.MAX_VALUE));
            } catch (IllegalBlockSizeException e) {
                e.printStackTrace();
            } catch (BadPaddingException e) {
                e.printStackTrace();
            }
        else
            dataOut.println(pckt.toString());
        return true;
    }

    public StealthNetPacket recvPacket() throws IOException {
        StealthNetPacket pckt = null;
        String str = null;
        if (decrypt != null)
            try {
                str = new String(decrypt.doFinal(Base64.decode(dataIn.readLine())));
            } catch (IllegalBlockSizeException e) {
                e.printStackTrace();
            } catch (BadPaddingException e) {
                e.printStackTrace();
            } catch (Base64DecodingException e) {
                e.printStackTrace();
            }
        else
            str = dataIn.readLine();

        System.out.println("[received] " + str);

        pckt = new StealthNetPacket(str);
//        if (!authenticatePacket(pckt)) {
//            // exit the thread if user who sends the package can not be authenticated by the signSignature
//            System.err.println("use");
//            System.exit(1);
//        }
        if (isPacketExpired(pckt)) {
            // exit the thread if the packet has expired
            System.err.println("Data has expired.");
            System.exit(1);
        }
        if (!verifyPacket(pckt)) {
            // exit the thread if MAC can't be verified
            System.err.println("Data cannot be verified by attached MAC.");
            System.exit(1);
        }
        return pckt;
    }

    private boolean authenticatePacket(StealthNetPacket pckt) {
        return false;
    }

    private boolean isPacketExpired(StealthNetPacket pckt) {
        Date now = new Date();
        Date delay = new Date(pckt.timestamp.getTime() + TOLERANCE);
        return now.after(delay);
    }

    private boolean verifyPacket(StealthNetPacket pckt) {
        return hmac == null || pckt.mac == null || pckt.mac.length == 0
                || toString(hmac.doFinal(concatenate(pckt.command, pckt.data))).equals(toString(pckt.mac));
    }

    public boolean recvReady() throws IOException {
        /*
         * System.out.println("Connected: " + commsSocket.isConnected()); System.out.println("Closed: " +
         * commsSocket.isClosed()); System.out.println("InClosed: " + commsSocket.isInputShutdown());
         * System.out.println("OutClosed: " + commsSocket.isOutputShutdown());
         */
        return dataIn.ready();
    }

    /*
     * authenticate channel with Diffie Hellman Key exchange
     */
    private void initDHKey() throws Exception {
        DHParameterSpec dhParamSpec = null;

        if (dhPublicKey == null) {
            System.out.println("Creating Diffie-Hellman parameters (takes VERY long) ...");
            AlgorithmParameterGenerator paramGen = AlgorithmParameterGenerator.getInstance("DH");
            paramGen.init(512);
            AlgorithmParameters params = paramGen.generateParameters();
            dhParamSpec = (DHParameterSpec) params.getParameterSpec(DHParameterSpec.class);
        } else
            dhParamSpec = ((DHPublicKey) dhPublicKey).getParams();

        KeyPairGenerator keyPairGen = KeyPairGenerator.getInstance("DH");
        keyPairGen.initialize(dhParamSpec);
        dhKeyPair = keyPairGen.generateKeyPair();

        System.out.println("Initialization ...");
        keyAgree = KeyAgreement.getInstance("DH");
        keyAgree.init(dhKeyPair.getPrivate());

        byte[] pubKeyEnc = dhKeyPair.getPublic().getEncoded();
        sendPacket(StealthNetPacket.CMD_DH_KEY, pubKeyEnc);
    }

    /*
     * exchange of Diffie-Hellman public key
     */
    private void recvDHPublicKey() throws Exception {
        while (dhPublicKey == null) {
            if (recvReady()) {
                StealthNetPacket pkt = recvPacket();
                if (pkt.command == StealthNetPacket.CMD_DH_KEY) {
                    KeyFactory keyFac = KeyFactory.getInstance("DH");
                    X509EncodedKeySpec x509KeySpec = new X509EncodedKeySpec(pkt.data);
                    dhPublicKey = keyFac.generatePublic(x509KeySpec);
//                } else if (pkt.command == StealthNetPacket.CMD_RSA_KEY)
//                    rsaPublicKey = KeyFactory.getInstance("RSA").generatePublic(new X509EncodedKeySpec(pkt.data));
                }
                Thread.sleep(100);
            }
        }
    }

    /*
     * initialise two ciphers to encrypt and decrypt message data
     */
    private void initCipher() throws Exception {
        keyAgree.doPhase(dhPublicKey, true);
        SecretKey aesKey = keyAgree.generateSecret("AES");

        keyAgree.doPhase(dhPublicKey, true);
        SecretKey iv = keyAgree.generateSecret("AES");

        encrypt = Cipher.getInstance("AES/CBC/PKCS5Padding");
        decrypt = Cipher.getInstance("AES/CBC/PKCS5Padding");

        encrypt.init(Cipher.ENCRYPT_MODE, aesKey, new IvParameterSpec(iv.getEncoded(), 0, 16));
        decrypt.init(Cipher.DECRYPT_MODE, aesKey, new IvParameterSpec(iv.getEncoded(), 0, 16));
    }

    private void initMac() throws Exception {
        // Generate secret key for HMAC-SHA1 from key exchange
        keyAgree.doPhase(dhPublicKey, true);
        SecretKey macKey = keyAgree.generateSecret("AES");

        // Get instance of Mac object implementing HMAC-SHA1, and
        // initialize it with the above secret key
        hmac = Mac.getInstance("HmacSHA1");
        hmac.init(macKey);
    }

    private void initRSAKeyPair() throws Exception {
        File privateKeyStore = new File(self + PRIVATE_KEY_EXTENTION);
        File publicKeyStore = new File(self + PUBLIC_KEY_EXTENTION);

        if (privateKeyStore.exists() && verifyKeystore(privateKeyStore)) {
            System.out.println("public/private keystores exist and verified.");
            rsaPrivateKey = getRSAPrivateKey(privateKeyStore);
        } else {
            System.out.println("public/private keystores do not exist, generating a new pair.");
            privateKeyStore.delete();
            publicKeyStore.delete();

            KeyPairGenerator keyGen = KeyPairGenerator.getInstance("RSA");
            SecureRandom random = SecureRandom.getInstance("SHA1PRNG");
            random.setSeed(new Date().getTime());
            keyGen.initialize(1024, random);
            KeyPair rsaKeyPair = keyGen.generateKeyPair();

            rsaPrivateKey = rsaKeyPair.getPrivate();

            saveRSAPrivateKey(rsaPrivateKey, privateKeyStore);
            saveRSAPublicKey(rsaKeyPair.getPublic(), publicKeyStore);

            // need to send public key to server
//            sendPacket(StealthNetPacket.CMD_RSA_KEY, rsaKeyPair.getPublic().getEncoded());
        }

//        if (userId.equals(SERVERNAME)) {
//            getRSAPublicKey(new File(SERVERNAME + PUBLIC_KEY_EXTENTION));
//        } else if (self.equals(SERVERNAME)) {
//            File userKeystore = new File(userId + PUBLIC_KEY_EXTENTION);
//            if (userKeystore.exists()) {
//                getRSAPublicKey(new File(userId + PUBLIC_KEY_EXTENTION));
//            }
//            else
//                recvRSAPublicKey();
//        }
//        
    }
//
//    private void recvRSAPublicKey() throws Exception {
//        while (rsaPublicKey == null) {
//            if (recvReady()) {
//                StealthNetPacket pkt = recvPacket();
//                if (pkt.command == StealthNetPacket.CMD_RSA_KEY)
//                    rsaPublicKey = KeyFactory.getInstance("RSA").generatePublic(new X509EncodedKeySpec(pkt.data));
//
//                Thread.sleep(100);
//            }
//        }
//    }
//
//    private void requestRSAPublicKey(String userId) {
//        sendPacket(StealthNetPacket.CMD_RSA_KEY, userId);
//    }
//
//    private void sendRSAPublicKey(String userId) throws Exception {
//        File keystore = new File(userId + PUBLIC_KEY_EXTENTION);
//        if (keystore.exists()) {
//            PublicKey key = getRSAPublicKey(keystore);
//            sendPacket(StealthNetPacket.CMD_RSA_KEY, key.getEncoded());
//        } else {
//            // can't have no public keystore since the user's signed in
//        }
//    }

    // Decrypt RSA private key and deserialise to PrivateKey object
    private PrivateKey getRSAPrivateKey(File keystore) throws Exception {
        // Generate the secret key specs
        PBEParameterSpec pbeParamSpec = new PBEParameterSpec(SALT, 20);
        PBEKeySpec pbeKeySpec = new PBEKeySpec(FILE_AES_KEY.toCharArray());
        SecretKeyFactory keyFactory = SecretKeyFactory.getInstance("PBEWithSHA1AndDESede");
        SecretKey pbeKey = keyFactory.generateSecret(pbeKeySpec);

        // Instantiate the cipher
        Cipher cipher = Cipher.getInstance("PBEWithSHA1AndDESede/CBC/PKCS5Padding");
        cipher.init(Cipher.DECRYPT_MODE, pbeKey, pbeParamSpec);

        FileInputStream fis = new FileInputStream(keystore);
        CipherInputStream cis = new CipherInputStream(fis, cipher);
        ObjectInputStream ins = new ObjectInputStream(cis);

        PrivateKey key = (PrivateKey) ins.readObject();
        ins.close();
        return key;
    }

    // Deserialise RSA public key into PublicKey object
    private PublicKey getRSAPublicKey(File keystore) throws Exception {
        FileInputStream fis = new FileInputStream(keystore);
        ObjectInputStream ins = new ObjectInputStream(fis);
        PublicKey key = (PublicKey) ins.readObject();
        ins.close();
        return key;
    }

    // Serialise encrypted private key to file
    private void saveRSAPrivateKey(PrivateKey key, File keystore) throws Exception {
        // Generate the secret key specs
        PBEParameterSpec pbeParamSpec = new PBEParameterSpec(SALT, 20);
        PBEKeySpec pbeKeySpec = new PBEKeySpec(FILE_AES_KEY.toCharArray());
        SecretKeyFactory keyFactory = SecretKeyFactory.getInstance("PBEWithSHA1AndDESede");
        SecretKey pbeKey = keyFactory.generateSecret(pbeKeySpec);

        // Instantiate the cipher
        Cipher cipher = Cipher.getInstance("PBEWithSHA1AndDESede/CBC/PKCS5Padding");
        cipher.init(Cipher.ENCRYPT_MODE, pbeKey, pbeParamSpec);

        // output the encrypted private key to file
        FileOutputStream fos = new FileOutputStream(keystore);
        CipherOutputStream cos = new CipherOutputStream(fos, cipher);
        ObjectOutputStream oos = new ObjectOutputStream(cos);
        oos.writeObject(key);
        oos.close();

        // generate MAC for the encrypted keystore
        signKeytore(keystore);
    }

    // Serialise public key to file
    private void saveRSAPublicKey(PublicKey key, File keystore) throws Exception {
        FileOutputStream fos = new FileOutputStream(keystore);
        ObjectOutputStream oos = new ObjectOutputStream(fos);
        oos.writeObject(key);
        oos.close();
    }

    private boolean verifyKeystore(File keystore) throws Exception {
        SecretKeySpec key = new SecretKeySpec(FILE_MAC_KEY.getBytes(), "AES");

        // Get instance of Mac object implementing HMAC-SHA1, and
        // initialize it with the above secret key
        Mac mac = Mac.getInstance("HmacSHA1");
        mac.init(key);

        FileInputStream fis = new FileInputStream(keystore);
        mac.update(IOUtils.readFully(fis, Integer.MAX_VALUE, true));
        fis.close();

        File sha1 = new File(self + SHA1_MAC_EXTENTION);
        Scanner cin = new Scanner(sha1);

        String origMac = cin.next();
        String newMac = toString(mac.doFinal());

        return origMac.equals(newMac);
    }

    private void signKeytore(File keystore) throws Exception {
        SecretKeySpec key = new SecretKeySpec(FILE_MAC_KEY.getBytes(), "AES");

        // Get instance of Mac object implementing HMAC-SHA1, and
        // initialize it with the above secret key
        Mac mac = Mac.getInstance("HmacSHA1");
        mac.init(key);

        FileInputStream fis = new FileInputStream(keystore);
        mac.update(IOUtils.readFully(fis, Integer.MAX_VALUE, true));
        fis.close();

        File sha1 = new File(self + SHA1_MAC_EXTENTION);
        FileWriter fw = new FileWriter(sha1);

        fw.write(toString(mac.doFinal()));
        fw.close();
    }

    private void initSignature() throws Exception {
        signSignature = Signature.getInstance("SHA1withRSA");
        signSignature.initSign(rsaPrivateKey);

        verifySignature = Signature.getInstance("SHA1withRSA");
        verifySignature.initVerify(rsaPublicKey);
    }

    private byte[] concatenate(byte command, byte[] data) {
        byte[] concat = new byte[data.length + 1];
        concat[0] = command;
        System.arraycopy(data, 0, concat, 1, data.length);
        return concat;
    }

    private String toString(byte[] data) {
        String str;
        int i, lowByte, highByte;

        str = "";
        for (i = 0; i < data.length; i++) {
            highByte = (data[i] >= 0) ? data[i] : 256 + data[i];
            lowByte = highByte & 15;
            highByte /= 16;
            str += HEXTABLE[highByte];
            str += HEXTABLE[lowByte];
        }
        return str;
    }

    public static void main(String... args) throws Exception {

        KeyPairGenerator keyGen = KeyPairGenerator.getInstance("DSA");
        SecureRandom random = SecureRandom.getInstance("SHA1PRNG", "SUN");
        random.setSeed(new Date().getTime());
        keyGen.initialize(1024, random);
        KeyPair dsaKeyPair = keyGen.generateKeyPair();

        System.out.println("private key: " + dsaKeyPair.getPrivate());
        System.out.println("public key: " + dsaKeyPair.getPublic());

        AlgorithmParameterGenerator paramGen = AlgorithmParameterGenerator.getInstance("DH");
        paramGen.init(512);
        AlgorithmParameters params = paramGen.generateParameters();
        DHParameterSpec dhParamSpec = (DHParameterSpec) params.getParameterSpec(DHParameterSpec.class);
        KeyPairGenerator keyPairGen = KeyPairGenerator.getInstance("DH");
        keyPairGen.initialize(dhParamSpec);
        KeyPair dhKeyPair = keyPairGen.generateKeyPair();

        System.out.println("private key: " + dhKeyPair.getPrivate());
        System.out.println("public key: " + dhKeyPair.getPublic());

    }
}

/******************************************************************************
 * END OF FILE: StealthNetComms.java
 *****************************************************************************/

